__all__=["vec_to_enu","solve_plane_pt","solve_plane_pt_new"]

import itertools

import cv2
import numpy as np

from bs_imu.bs_imu import rot_vec_to_matrix
from .bs_cfg_real_5m import *

def vec_to_enu(t_vec,r_vec,R_wt):
    # PnP解算出来的R是靶标坐标系到相机系中的坐标
    # 逆变换求相机系的原点在靶标坐标系中的位置
    # 然后再从靶标坐标系变换到世界系
    R_ct = rot_vec_to_matrix(r_vec)
    return -np.dot(R_wt,R_ct.T.dot(t_vec))


def solve_plane_pt(plane_px_ptL,plane_real_ptL,cam_K,distort_coeffs):

    '''
    传入特征点列表，通过飞机特征点解算相对位姿
    '''
    result_dict = {
        "rpe": np.nan,
        "r_vec": np.array([np.nan]*3),
        "t_vec": np.array([np.nan]*3),
        "plane_px_calcL": np.array([]),
        "img_valid": False,
    }

    pt_num = len(plane_px_ptL)
    # 特征点不足以解算
    if pt_num < 5:
        return result_dict

    data_dictL = []
    # 从左至右排序
    plane_px_ptL_sorted = sorted(plane_px_ptL,key=lambda x:x[0])
    
    # 特征点大于5个以上
    for combineL in itertools.combinations(range(len(plane_px_ptL)),5):
        plane_px_calcL = np.array([
            plane_px_ptL_sorted[i]
            for i in combineL
        ])
        # retval,r_vec,t_vec = cv2.solvePnP(plane_real_ptL, plane_px_calcL, cam_K, distCoeffs=None, flags=cv2.SOLVEPNP_EPNP)
        retval,r_vec,t_vec = cv2.solvePnP(plane_real_ptL, plane_px_calcL, cam_K, distCoeffs=distort_coeffs, flags=cv2.SOLVEPNP_SQPNP)
        # 计算重投影误差
        reproj_errL = []
        for plane_real_pt,plane_px_calc in zip(plane_real_ptL,plane_px_calcL):
            plane_proj_pt,_ = cv2.projectPoints(plane_real_pt, r_vec, t_vec, cam_K, distCoeffs=None)
            #normalize based on the original size 
            plane_px_err = np.linalg.norm(plane_proj_pt.flatten() - plane_px_calc )/np.linalg.norm(plane_px_calcL[0,:]-plane_px_calcL[4,:])  #position error
            reproj_errL.append(plane_px_err)
        reproj_err = np.max(reproj_errL)
        data_dict = {
            "rpe": reproj_err,
            "r_vec": r_vec.flatten(),
            "t_vec": t_vec.flatten(),
            "plane_px_calcL": plane_px_calcL,
        }
        data_dictL.append(data_dict)

    
    # 取重投影误差最小的
    data_dictL_sorted = sorted(data_dictL,key=lambda x: x["rpe"],reverse=False)
    data_dict = data_dictL_sorted[0]
   
    # 重投影误差过大
    if data_dict["rpe"] < 0.1:
        result_dict = data_dict
        result_dict["img_valid"] = True
    return result_dict


def solve_plane_pt_new(plane_px_ptL,plane_real_ptL,cam_K,distort_coeffs,plane_area_ptL):

    '''
    传入特征点列表，通过飞机特征点解算相对位姿
    '''
    result_dict = {
        "rpe": np.nan,
        "r_vec": np.array([np.nan]*3),
        "t_vec": np.array([np.nan]*3),
        "plane_px_calcL": np.array([]),
        "img_valid": False,
    }

    pt_num = len(plane_px_ptL)
    # 特征点不足以解算
    if pt_num < 5:
        return result_dict

    data_dictL = []
    # 从左至右排序
    plane_px_ptL_sorted = sorted(plane_px_ptL,key=lambda x:x[0])
    temp = list(zip(plane_px_ptL, plane_area_ptL))
    sorted_temp = sorted(temp, key=lambda x: x[0][0])
    plane_px_ptL_sorted, plane_area_ptL_sorted = zip(*sorted_temp)
    # 特征点大于5个以上
    for combineL in itertools.combinations(range(len(plane_px_ptL_sorted)),5):
        plane_px_calcL = np.array([
            plane_px_ptL_sorted[i]
            for i in combineL
        ])
        plane_area_calcL = np.array([
            plane_area_ptL_sorted[i]
            for i in combineL
        ])
        # retval,r_vec,t_vec = cv2.solvePnP(plane_real_ptL, plane_px_calcL, cam_K, distCoeffs=None, flags=cv2.SOLVEPNP_EPNP)
        retval,r_vec,t_vec = cv2.solvePnP(plane_real_ptL, plane_px_calcL, cam_K, distCoeffs=distort_coeffs, flags=cv2.SOLVEPNP_SQPNP)
        if r_vec is not None and t_vec is not None:
            if not(any(x is None for x in t_vec)) and not(any(x is None for x in r_vec)):
                # 计算重投影误差
                reproj_errL = []
                for plane_real_pt,plane_px_calc,plane_area_pt in zip(plane_real_ptL,plane_px_calcL,plane_area_ptL):
                    plane_proj_pt,_ = cv2.projectPoints(plane_real_pt, r_vec, t_vec, cam_K, distCoeffs=distort_coeffs)
                    #normalize based on the original size 
                    plane_px_err = np.linalg.norm((plane_proj_pt.flatten() - plane_px_calc )/(plane_area_pt))/(np.linalg.norm(plane_px_calcL[0,:]-plane_px_calcL[4,:]))  #position error
                    reproj_errL.append(plane_px_err)
                reproj_err = np.mean(reproj_errL)
                data_dict = {
                    "rpe": reproj_err,
                    "r_vec": r_vec.flatten(),
                    "t_vec": t_vec.flatten(),
                    "plane_px_calcL": plane_px_calcL,
                }
                data_dictL.append(data_dict)
    
    # 取重投影误差最小的
    data_dict = {
        "rpe": np.nan,
        "r_vec": [np.nan]*3,
        "t_vec": [np.nan]*3,
        "plane_px_calcL": [np.nan]*5,
    }
    data_dictL_sorted = sorted(data_dictL,key=lambda x: x["rpe"])
    i = 0
    l = len(data_dictL_sorted)
    for i in range(0,l):
        data_dict = data_dictL_sorted[i]
        plane_pts = data_dict["plane_px_calcL"]
        ar1 = np.linalg.norm(plane_pts[0,:]-plane_pts[1,:])
        ar2 = np.linalg.norm(plane_pts[1,:]-plane_pts[2,:])
        ar3 = np.linalg.norm(plane_pts[2,:]-plane_pts[3,:])
        ar4 = np.linalg.norm(plane_pts[3,:]-plane_pts[4,:])
        ar5 = np.linalg.norm(plane_pts[0,:]-plane_pts[2,:])
        ar6 = np.linalg.norm(plane_pts[2,:]-plane_pts[4,:])
        ar7 = np.linalg.norm(plane_pts[1,:]-plane_pts[2,:])
        ar8 = np.linalg.norm(plane_pts[2,:]-plane_pts[3,:])
        plane_shape_val = np.array([ar1/ar2,ar2/ar3,ar3/ar4,ar5/ar6,ar7/ar8])
        plane_shape_err = np.linalg.norm(plane_shape_val-plane_real_val)
        if plane_shape_err < 3:
            break

    # print("plane",data_dict["rpe"],plane_shape_err,data_dict["t_vec"])
    # 重投影误差过大
    if not np.isnan(data_dict["rpe"]) and data_dict["rpe"] < 0.001:
        result_dict = data_dict
        result_dict["img_valid"] = True
    return result_dict

